'''
最速下降方法
Author：Zhekai Zhang
Time: 2024/12/5 23:20
Email: zhangzhk26@mail2.sysu.edu.cn
'''

import sympy as sp
import numpy as np
import matplotlib.pyplot as plt


def steepest_descent(func_expr, t, epsilon=1e-6, method="定步长", alpha=0.1, a=0, b=1, max_iter=1):
    """
    最速下降法函数
    参数：
    - func_expr：目标函数的表达式（例如: "(1 - x)**2 + 100*(y - x**2)**2"）
    - x0, y0：起始点的坐标
    - epsilon：收敛精度
    - method：步长选择方法 ("定步长" 或 "变步长")
    - alpha：定步长法的步长
    - a, b：黄金分割法的搜索区间
    - max_iter：最大迭代次数
    返回：
    - 最优解的坐标和迭代过程的点列表
    """

    x0 = t[0]
    y0 = t[1]

    # 定义符号变量
    x, y = sp.symbols('x y')

    # 将字符串转换为 sympy 表达式
    func = sp.sympify(func_expr)

    # 将符号函数转换为可计算的函数
    func_lambda = sp.lambdify((x, y), func, modules=['numpy'])

    # 计算梯度函数
    grad_func = [sp.diff(func, var) for var in (x, y)]
    grad_lambda = [sp.lambdify((x, y), g, modules=['numpy']) for g in grad_func]

    # 设置初始点
    point = np.array([x0, y0], dtype=float)

    # # 设置绘图
    # fig, ax = plt.subplots(figsize=(8, 6))

    # # 绘制目标函数的等高线图
    # x_vals = np.linspace(-20, 20, 400)
    # y_vals = np.linspace(-10, 30, 400)
    # X, Y = np.meshgrid(x_vals, y_vals)
    # Z = func_lambda(X, Y)

    # ax.contour(X, Y, Z, levels=np.logspace(-0.5, 3, 35), cmap='jet')
    # ax.set_title('contour')
    # ax.set_xlabel('x')
    # ax.set_ylabel('y')

    points = []  # List to store the points (x, y, func_value)

    # 定步长方法
    if method == "定步长":
        with open("定步长迭代过程.txt", "w", encoding="utf-8") as file:
            file.write(f"目标函数：{func_expr}\n")
            file.write(f"初始点：({x0}, {y0})\n")
            file.write(f"收敛精度：{epsilon}\n")
            file.write(f"步长 alpha：{alpha}\n")

            # 开始迭代
            for i in range(1, max_iter + 1):
                grad_value = np.array([g(point[0], point[1]) for g in grad_lambda], dtype=float)
                grad_norm = np.linalg.norm(grad_value)

                if grad_norm < epsilon:
                    file.write("达到了目标精度\n")
                    break

                # 更新点
                point = point - alpha * grad_value
                func_value = func_lambda(point[0], point[1])
                file.write(f"迭代次数：{i}\n")
                file.write(f"当前点：({point[0]}, {point[1]})\n")
                file.write(f"梯度：{grad_value}\n")
                file.write(f"梯度范数：{grad_norm}\n")
                file.write(f"函数值：{func_value}\n\n")

                # 保存每次迭代的 (x, y, function value)
                points.append((point[0], point[1], func_value))

                # if i % 100 == 0:
                #     ax.plot([p[0] for p in points], [p[1] for p in points], 'ro-', markersize=3)
                #     plt.pause(0.01)

            else:
                file.write("达到最大迭代次数\n")

        # ax.plot([p[0] for p in points], [p[1] for p in points], 'ro-', markersize=3, label='定步长迭代点')
        # ax.legend()
        # plt.show()

    # 变步长方法：黄金分割法和Armijo法
    elif method == "变步长":
        def golden_section_search(func, point, direction, a, b, tol=1e-5):
            phi = (1 + np.sqrt(5)) / 2
            resphi = 2 - phi

            x1 = a + resphi * (b - a)
            x2 = b - resphi * (b - a)
            f1 = func(*(point + x1 * direction))
            f2 = func(*(point + x2 * direction))

            while abs(b - a) > tol:
                if f1 < f2:
                    b = x2
                    x2 = x1
                    f2 = f1
                    x1 = a + resphi * (b - a)
                    f1 = func(*(point + x1 * direction))
                else:
                    a = x1
                    x1 = x2
                    f1 = f2
                    x2 = b - resphi * (b - a)
                    f2 = func(*(point + x2 * direction))

            return (a + b) / 2

        def armijo_search(func, grad_funcs, point, direction, alpha=1, rho=0.8, c=1e-4):
            func_value = func(*point)
            grad_value = np.array([g(*point) for g in grad_funcs])
            while func(*(point + alpha * direction)) > func_value + c * alpha * np.dot(grad_value, direction):
                alpha *= rho
            return alpha

        max_iter = 1

        # 黄金分割法
        point_golden = point.copy()
        points_golden = []
        with open("黄金分割法迭代过程.txt", "w", encoding="utf-8") as file_golden:
            file_golden.write(f"目标函数：{func_expr}\n")
            file_golden.write(f"初始点：({x0}, {y0})\n")
            file_golden.write(f"收敛精度：{epsilon}\n")

            for i in range(1, max_iter + 1):
                grad_value = np.array([g(point_golden[0], point_golden[1]) for g in grad_lambda], dtype=float)
                grad_norm = np.linalg.norm(grad_value)

                if grad_norm < epsilon:
                    file_golden.write("达到了目标精度（黄金分割法）\n")
                    break

                direction = -grad_value
                alpha = golden_section_search(func_lambda, point_golden, direction, a, b)
                point_golden = point_golden + alpha * direction
                func_value = func_lambda(point_golden[0], point_golden[1])
                file_golden.write(f"迭代次数：{i}\n")
                file_golden.write(f"当前点：({point_golden[0]}, {point_golden[1]})\n")
                file_golden.write(f"步长 alpha：{alpha}\n")
                file_golden.write(f"梯度：{grad_value}\n")
                file_golden.write(f"梯度范数：{grad_norm}\n")
                file_golden.write(f"函数值：{func_value}\n\n")

                # 保存每次迭代的 (x, y, function value)
                points_golden.append((point_golden[0], point_golden[1], func_value))

                # if i % 100 == 0:
                #     ax.plot([p[0] for p in points_golden], [p[1] for p in points_golden], 'ro-', markersize=3)
                #     plt.pause(0.01)

            else:
                file_golden.write("达到最大迭代次数（黄金分割法）\n")

        # ax.plot([p[0] for p in points_golden], [p[1] for p in points_golden], 'ro-', markersize=3,
        #         label='黄金分割法迭代点')
        # ax.legend()
        # plt.show()

    else:
        print("无效的选择，请选择“定步长”或“变步长”")
        return

    return points


# # 调用示例：
# points = steepest_descent("(1 - x)**2 + 100*(y - x**2)**2", 0, 0)
# print("迭代点：")
# for point in points:
#     print(f"x = {point[0]}, y = {point[1]}, function value = {point[2]}")
